import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np
import os.path as osp
import math, random


def meshgrid_2v(x, y, mesh_fn="mine" or "numpy"):
    """Return coordinate matrices from 2 coordinate vectors `x` and `y`."""

    assert mesh_fn in ["mine", "numpy"], 'choose either "mine" or "numpy" meshgrid'

    if mesh_fn == "mine":
        # Note that broadcasting happens on the highest dimension
        X = np.broadcast_to(x, (len(y), len(x)))
        Y = np.broadcast_to(y, (len(x), len(y))).transpose()

    if mesh_fn == "numpy":
        X, Y = np.meshgrid(x, y)

    return X, Y


def asin_ux_vy_1d(x, y, a=1, u=1, v=1):
    """Calculate `f(x, y) = a * sin(ux + vy)` at the given point `(x, y)`."""

    return a * math.sin(u * x + v * y)


def asin_ux_vy_2d(m, n, a=1, u=1, v=1, mesh_fn="mine"):
    """Calculate `f(X, Y) = a * sin(uX + vY)` where `X` and `Y` the intersections
    of line group `x = 0, ..., m - 1` and line group `y = 0, ..., n - 1`.

    This function utilize the broadcast operation to complete the task.
    """

    assert m > 1 and n > 1, '"m" and "n" should be non-negative integers'

    x = np.arange(0, m, 1)  # x = 0, 1, ..., m - 1
    y = np.arange(0, n, 1)  # y = 0, 1, ..., n - 1
    X, Y = meshgrid_2v(x, y, mesh_fn)

    return a * np.sin(u * X + v * Y)


def plot_ax(ax, X, Y, Z, *, title=None):
    """Plot the 3D surface in the given axes. The style of each subplot,
    such as linewidth, colormap and ticks should be configured here.
    """

    ax.plot_surface(X, Y, Z, cmap=cm.viridis, linewidths=0)
    if title:
        ax.set_title(title, fontsize=8)
    ax.set_xlabel("x")
    ax.set_ylabel("y")


def plot_figure(fig, u, v, *, m=256, n=256, a=32, mesh_fn="mine"):
    """Plot the 3D surface of the calculated plane for function `asin_ux_vy_2d`,
    given the values of `u` and the values of `v`, in the given figure.

    Both `u` and `v` can be either a scalar or an array of scalars.
    """

    u = u if isinstance(u, (list, tuple)) else [u]
    v = v if isinstance(v, (list, tuple)) else [v]

    x = np.arange(0, m, 1)  # x = 0, 1, ..., m - 1
    y = np.arange(0, n, 1)  # y = 0, 1, ..., n - 1
    X, Y = meshgrid_2v(x, y, mesh_fn)

    for udx, u_value in enumerate(u):
        for vdx, v_value in enumerate(v):
            Z = asin_ux_vy_2d(m, n, a, u_value, v_value, mesh_fn)
            ax = fig.add_subplot(
                len(u), len(v), len(v) * udx + vdx + 1, projection="3d"
            )
            plot_ax(ax, X, Y, Z, title=f"(u = {u_value}, v = {v_value})")

    fig.suptitle(f"f(x, y) = a * sin(ux + vy), m={m}, n={n}, a={a}", y=0.96)


def unit_test_problem_01(n_tests=6):
    """Problem 01: write a function to calculate `z = f(x, y) = a * sin(ux + vy).`"""

    problem = (
        "[ a = {a:.2f}, u = {u:.2f}, v = {v:.2f} ] => f({x:.2f}, {y:.2f}) = {r:.2f}"
    )

    print(
        f"Problem 01: write a function to calculate `z = f(x, y) = a * sin(ux + vy).`"
    )
    for tdx in range(n_tests):
        a, u, v, x, y = [(random.random() - 0.5) * 5.0 for _ in range(5)]
        print(problem.format(a=a, u=u, v=v, x=x, y=y, r=asin_ux_vy_1d(x, y, a, u, v)))


def unit_test_problem_02(u_values=0.2, v_values=0.4):
    """Problem 02: given m = n = 256, a = 32, plot the surface of `f(x, y)` with different u and v."""

    print(
        f"Problem 02: given m = n = 256, a = 32, plot the surface of `f(x, y)` with different u and v."
    )

    u_values = u_values if isinstance(u_values, (list, tuple)) else [u_values]
    v_values = v_values if isinstance(v_values, (list, tuple)) else [v_values]

    fig = plt.figure(figsize=(len(v_values) * 3.0 + 1.0, len(u_values) * 3.0 + 1.0))
    plot_figure(fig, u_values, v_values, m=256, n=256, a=32)
    fig_path = osp.join(osp.dirname(__file__), "asin_ux_vy_2d.png")
    plt.savefig(fig_path)

    print(f'[ u = {u_values}, v = {v_values}] => "{fig_path}"')


if __name__ == "__main__":
    unit_test_problem_01(n_tests=4)
    unit_test_problem_02([0.02, 0.04], [0.03, 0.04, 0.05])
